#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time         : 17-11-22 下午2:26
# @Author       : xiaohan
# @File         : Regression.py
# @desc         : 

import torch
from torch.autograd import Variable  # 包入数据
import torch.nn.functional as F  # 激励函数
import matplotlib.pyplot as plt  # 画图模块

# unsqueeze把一维的数据转换成2维
x = torch.unsqueeze(torch.linspace(-1, 1, 100), dim=1)  # x data (tensor) ,shape=(100,2)
y = x.pow(2) + 0.2 * torch.rand(x.size())

x, y = Variable(x), Variable(y)


# 散点图
# plt.scatter(x.data.numpy(), y.data.numpy())
#
# plt.show()


# net的主模块  torch.nn.Module
class Net(torch.nn.Module):
    # 搭建初始化这一层所需要的细信息
    def __init__(self, in_feature, n_hidden, n_output):
        # 调用nn模块
        super(Net, self).__init__()

        # 建立自己层的模块 hidden 层   层信息都是这个模块中的属性，这个属性中的内容就是层信息 torch.nn.Linear()
        self.hidden = torch.nn.Linear(in_feature, n_hidden)
        self.predict = torch.nn.Linear(n_hidden, n_output)

    # 前向传递的过程
    def forward(self, x):
        # 激励函数激活
        x = F.relu(self.hidden(x))
        x = self.predict(x)

        return x


net = Net(1, 10, 1)
print(net)

plt.ion()
plt.show()

# 优化器优化网络参数  lr学习效率
optimizer = torch.optim.SGD(net.parameters(), lr=0.5)

# 计算误差(均方差) --->处理回归问题  分类问题用其他方式计算分类误差
loss_func = torch.nn.MSELoss()

for t in range(100):
    prediction = net(x)
    # prediction 预测值  有真实值
    loss = loss_func(prediction, y)

    # 梯度圈全部将为0
    optimizer.zero_grad()

    # fan
    loss.backward()
    optimizer.step()


    if t %5 ==0:
        plt.cla()
        plt.scatter(x.data.numpy(),y.data.numpy())
        plt.plot(x.data.numpy(),prediction.data.numpy(),'r-',lw=5)
        plt.text(0.5,0,'Loss=%.4f' % loss.data[0],fontdict={'size':20,'color':'red'})

        plt.pause(0.1)
plt.ioff()
plt.show()
